import argparse
import json
import math
import os, cv2

import numpy as np
from tqdm import tqdm


parser = argparse.ArgumentParser('rec')
parser.add_argument("--basepath", type=str, default='E:/data/BaiduNetdiskDownload/hphm_02_0618_6000')
parser.add_argument("--imgsavepath", type=str, default='plate_rec_6000')
parser.add_argument("--txtsavepath", type=str, default='E:/data/BaiduNetdiskDownload/hphm_02_0618_6000/train_hphm_rec.txt')
parser.add_argument("--det_txt", type=str, default='E:/data/BaiduNetdiskDownload/hphm_02_0618_6000/train_hphm.txt')
args = parser.parse_args()

if not os.path.exists(os.path.join(args.basepath,args.imgsavepath)):
    os.mkdir(os.path.join(args.basepath,args.imgsavepath))
count = 0



def rotate(ps, m):
    pts = np.float32(ps).reshape([-1, 2])  # 要映射的点
    pts = np.hstack([pts, np.ones([len(pts), 1])]).T
    target_point = np.dot(m, pts)
    target_point = [[int(target_point[0][x]),int(target_point[1][x])] for x in range(len(target_point[0]))]
    return target_point

def rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
    h,w,c = img.shape
    M = cv2.getRotationMatrix2D((center_x,center_y), angle, resize_rate)
    res_img = cv2.warpAffine(img, M, (w, h))
    out_points = rotate(points, M)
    return res_img,out_points

def rotate_point(points,angle,center_x,center_y,resize_rate=1.0):
    M = cv2.getRotationMatrix2D((center_x, center_y), angle, resize_rate)
    out_points = rotate(points, M)
    return  out_points

def distance(point1,point2):
    return int(math.sqrt((point1[0]-point2[0])**2+(point1[1]-point2[1])**2))

def paixu(p):
    p = sorted(p, key=lambda point: point[1])
    sorts = []
    if p[0][0]<p[1][0]:
        sorts.append(p[0])
        sorts.append(p[1])
    else:
        sorts.append(p[1])
        sorts.append(p[0])
    if p[2][0]<p[3][0]:
        sorts.append(p[3])
        sorts.append(p[2])
    else:
        sorts.append(p[2])
        sorts.append(p[3])
    return sorts


def imggcut(image,y0,y1,x0,x1):
    if y0<0:
        y0 = 0
    if y1<0:
        y1 = 0
    if x0< 0:
        x0=0
    if x1 <0:
        x1=0
    return image[y0:y1,x0:x1]

def get_center(p1,p2,p3,p4):
    cx = int((p1[0]+p2[0]+p3[0]+p4[0])/4.0)
    cy = int((p1[1]+p2[1]+p3[1]+p4[1])/4.0)
    return [cx, cy]


def process(img, sorts):
    sp = img.shape
    h = sp[0]
    w = sp[1]

    w_1 = sorts[1][0]-sorts[0][0]
    h_1 = sorts[2][1]-sorts[1][1]


    angel = math.acos((sorts[1][0] - sorts[0][0]) / distance(sorts[0], sorts[1]))

    #print("angel",angel)
    # rotate_img_and_point(img,points,angle,center_x,center_y,resize_rate=1.0):
    if sorts[1][1] > sorts[0][1]:
        src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                          (angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)
    else:
        src_img, p = rotate_img_and_point(img, get_center(sorts[0], sorts[1], sorts[2], sorts[3]),
                                          (-1)*(angel / math.pi * 180), sp[1] / 2.0, sp[0] / 2.0)


    #
    src_img_w =src_img.shape[1]
    src_img_h = src_img.shape[0]
    y0 = int(p[0][1] - h_1*0.5)
    x0 = int(p[0][0] - w_1*0.5)
    y1 = int(p[0][1] + h_1 * 0.5)
    x1 = int(p[0][0] + w_1 * 0.5)

    #print(y0,y1, x0,x1)
    if x0 < 0:
        x0 = 0
    if y0 < 0:
        y0 = 0
    if x1 > src_img_w:
        x1 = src_img_w
    if y1 > src_img_h:
        y1 = src_img_h
    src_img = src_img[y0:y1, x0:x1]

    return src_img


f = open(args.txtsavepath,"w",encoding='utf-8')
count = 0

for i in tqdm(open(args.det_txt,"r",encoding='utf-8').readlines()):
    img_append_path, anonation = i.split('\t')
    json_q = json.loads(anonation)
    json_ = json_q[0]
    name = os.path.basename(img_append_path)
    img = cv2.imread(os.path.join(args.basepath,img_append_path))
    if img is None:
        print(name)
    else:
        if len(json_q) ==1:
            count = count + 1
            img_cut = process(img,json_['points'])
            cv2.imwrite(os.path.join(args.basepath,args.imgsavepath,'%06d.jpg'%count),img_cut)
            f.write(os.path.join(args.imgsavepath,'%06d.jpg'%count)+"\t"+json_["transcription"]+"\n")


f.close()


